import onnx
import onnx_graphsurgeon as gs
import argparse
import os
from typing import List, Optional
import numpy as np


def remove_nodes_with_graphsurgeon(model_path: str, output_path: str, nodes_to_remove: List[str]):

    print(f"加载ONNX模型: {model_path}")
    
    # 加载ONNX模型并转换为GraphSurgeon图
    graph = gs.import_onnx(onnx.load(model_path))
    
    print("模型加载完成，开始分析图结构...")
    
    # 获取所有节点的名称
    all_node_names = {node.name for node in graph.nodes if node.name}
    print(f"模型中总共有 {len(graph.nodes)} 个节点")
    
    # 检查要删除的节点是否存在
    existing_nodes_to_remove = []
    missing_nodes = []
    
    for node_name in nodes_to_remove:
        if node_name in all_node_names:
            existing_nodes_to_remove.append(node_name)
        else:
            missing_nodes.append(node_name)
    
    if missing_nodes:
        print(f"警告: 以下节点未找到: {missing_nodes}")
    
    if not existing_nodes_to_remove:
        print("没有找到需要删除的节点")
        return []
    
    print(f"准备删除以下节点: {existing_nodes_to_remove}")
    
    # 删除节点
    removed_nodes = []
    nodes_to_remove_objects = []
    
    # 找到要删除的节点对象
    for node in graph.nodes:
        if node.name in existing_nodes_to_remove:
            nodes_to_remove_objects.append(node)
            removed_nodes.append(node.name)
            print(f"  找到节点: {node.name} (操作类型: {node.op})")
    
    # 分析节点连接关系并进行智能重连
    print("分析节点连接关系...")
    for node in nodes_to_remove_objects:
        print(f"  处理节点: {node.name}")
        print(f"    输入: {[inp.name if hasattr(inp, 'name') else str(inp) for inp in node.inputs]}")
        print(f"    输出: {[out.name if hasattr(out, 'name') else str(out) for out in node.outputs]}")
        
        
        # 对于有多个输入输出的节点，我们需要手动重连
        if len(node.inputs) == 1 and len(node.outputs) == 1:
            # 简单的一对一连接
            input_tensor = node.inputs[0]
            output_tensor = node.outputs[0]
            
            # 找到所有使用这个输出的节点，并将它们的输入重定向到原始输入
            for consumer_node in graph.nodes:
                for i, inp in enumerate(consumer_node.inputs):
                    if inp is output_tensor:
                        consumer_node.inputs[i] = input_tensor
                        print(f"    重连: {consumer_node.name} 的输入 {i} 从 {output_tensor.name} 改为 {input_tensor.name}")
        
        elif len(node.inputs) == len(node.outputs):
            # 多对多连接，按索引对应
            for i, (input_tensor, output_tensor) in enumerate(zip(node.inputs, node.outputs)):
                for consumer_node in graph.nodes:
                    for j, inp in enumerate(consumer_node.inputs):
                        if inp is output_tensor:
                            consumer_node.inputs[j] = input_tensor
                            print(f"    重连: {consumer_node.name} 的输入 {j} 从 {output_tensor.name} 改为 {input_tensor.name}")
        else:
            print(f"    警告: 节点 {node.name} 的输入输出数量不匹配 (输入:{len(node.inputs)}, 输出:{len(node.outputs)})")
            print("    将尝试简单的第一个输入到所有输出的连接")
            if node.inputs:
                input_tensor = node.inputs[0]
                for output_tensor in node.outputs:
                    for consumer_node in graph.nodes:
                        for j, inp in enumerate(consumer_node.inputs):
                            if inp is output_tensor:
                                consumer_node.inputs[j] = input_tensor
                                print(f"    重连: {consumer_node.name} 的输入 {j} 从 {output_tensor.name} 改为 {input_tensor.name}")
    
    # 从图中移除节点
    print("从图中移除节点...")
    for node in nodes_to_remove_objects:
        graph.nodes.remove(node)
        print(f"  已移除节点: {node.name}")
    
    # 清理图：移除未使用的张量和进行优化
    print("清理和优化图...")
    graph.cleanup()
    
    # 尝试进行形状推理
    try:
        print("进行形状推理...")
        # 导出为ONNX进行形状推理
        temp_model = gs.export_onnx(graph)
        temp_model = onnx.shape_inference.infer_shapes(temp_model)
        # 重新导入以获得形状信息
        graph = gs.import_onnx(temp_model)
        print("形状推理完成")
    except Exception as e:
        print(f"形状推理失败: {e}")
        print("继续保存模型...")
    
    # 导出修改后的模型
    print(f"保存修改后的模型到: {output_path}")
    modified_model = gs.export_onnx(graph)
    
    # 验证模型
    try:
        onnx.checker.check_model(modified_model)
        print("模型验证通过")
    except Exception as e:
        print(f"模型验证失败: {e}")
        print("警告: 生成的模型可能需要进一步检查")
    
    # 保存模型
    onnx.save(modified_model, output_path)
    print(f"成功删除了 {len(removed_nodes)} 个节点: {removed_nodes}")
    
    return removed_nodes

def main():
    """主函数"""
    # 设置默认参数
    default_model_path = "your_model.onnx"
    default_nodes = ["your_node_name"]   # 需要移除的节点名称列表

    
    parser = argparse.ArgumentParser(description="使用GraphSurgeon删除ONNX模型中的指定节点")
    parser.add_argument("--model_path", type=str, default=default_model_path,
                        help="输入ONNX模型路径")
    parser.add_argument("--output_path", type=str, default=None,
                        help="输出ONNX模型路径，默认在原文件名后添加_removed_nodes")
    parser.add_argument("--nodes", nargs="+", default=default_nodes,
                        help="要删除的节点名称列表")
    
    args = parser.parse_args()
    
    # 检查输入文件是否存在
    if not os.path.exists(args.model_path):
        print(f"错误: 模型文件不存在: {args.model_path}")
        return
    
    # 生成输出路径
    if args.output_path is None:
        base_name = os.path.splitext(args.model_path)[0]
        args.output_path = f"{base_name}_removed_nodes.onnx"
    
    print("=" * 60)
    print("ONNX节点删除工具 (基于GraphSurgeon)")
    print("=" * 60)
    print(f"输入模型: {args.model_path}")
    print(f"输出模型: {args.output_path}")
    print(f"要删除的节点: {args.nodes}")
    print("-" * 60)
    
    try:
        # 删除节点
        removed_nodes = remove_nodes_with_graphsurgeon(args.model_path, args.output_path, args.nodes)
        
        print("-" * 60)
        print("操作完成!")
        if removed_nodes:
            print(f"已删除的节点: {removed_nodes}")
            print(f"修改后的模型已保存到: {args.output_path}")
        else:
            print("没有删除任何节点")
        print("=" * 60)
            
    except Exception as e:
        print(f"错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
